package com.senlin;

import java.io.IOException;
import java.io.InputStream;
import java.nio.file.*;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class MemoryDownloadHandler extends DefaultDownloadHandler {

    private final ConcurrentHashMap<Long, byte[]> dataCache = new ConcurrentHashMap<>();

    @Override
    public void data(InputStream inputStream, DownloadTask downloadTask, ShardingTask shardingTask) throws Exception {
        byte[] bytes = inputStream.readAllBytes();
        dataCache.put(shardingTask.getOffset(), bytes);
        System.out.println("[下载完成@" + shardingTask.getId() + "#" + shardingTask.getOffset() + "]  - " + shardingTask);
    }

    @Override
    public void completed(DownloadTask task) {
        final Path path = Paths.get(task.getLocation());
        dataCache.entrySet().stream()
                .sorted(Map.Entry.comparingByKey())
                .map(Map.Entry::getValue)
                .forEach(bytes -> {
                    try {
                        Files.write(path, bytes, StandardOpenOption.CREATE, StandardOpenOption.APPEND);
                    } catch (IOException e) {
                        e.printStackTrace();
                        // ignore
                    }
                });
        dataCache.clear();
        super.completed(task);
    }
}
